import torch
import torch.nn as nn
import os
from transformers import BertModel, BertTokenizer
import torch.nn.functional as F
# from pytorch_pretrained import BertModel, BertTokenizer


class Config(object):

    """配置参数"""
    def __init__(self, dataset):
        self.model_name = 'ERNIE'
        self.train_path = '/data/sfq/oco/train.csv'                                   # 训练集
        self.test_path = '/data/sfq/oco/test_new.csv'                                 # 测试集
        self.stop_words = '/data/sfq/oco/stopwords.txt' 
        self.save_path = '/data/sfq/oco/model.pt'                                # 模型训练结果
        self.log_path = '/data/sfq/oco/'
        self.device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')   # 设备

        self.require_improvement = 1000                                 # 若超过1000batch效果还没提升，则提前结束训练
        self.num_classes = 2                                            # 类别数
        self.num_epochs = 10                                            # epoch数
        self.batch_size = 512                                           # mini-batch大小
        self.pad_size = 128                                             # 每句话处理成的长度(短填长切) 32
        self.learning_rate = 5e-5                                       # 学习率
        self.bert_path = '/data/pretrain_models/bert_cn/'
        self.hidden_size = 768
        self.filter_sizes = (2, 3, 4)                                   # 卷积核尺寸
        self.num_filters = 8                                            # 卷积核数量(channels数)
        self.dropout = 0.1
        self.rnn_hidden = 768
        self.num_layers = 2
        self.gpu_id = 1
        self.n_gram_vocab = 250499
        self.encoder_lr = 5e-4
        self.head_lr = 5e-3
        self.fold = 5
        self.tokenizer = None


class Model(nn.Module):

    def __init__(self, config):
        super(Model, self).__init__()
        self.bert = BertModel.from_pretrained(config.bert_path)
        self.lstm = nn.LSTM(config.hidden_size, config.rnn_hidden, config.num_layers,
                            bidirectional=True, batch_first=True, dropout=config.dropout)
        self.dropout = nn.Dropout(config.dropout)
        self.fc_rnn = nn.Linear(config.rnn_hidden * 2, config.num_classes)

    def forward(self, x):
        context = x[0] 
        mask = x[2] 
        output = self.bert(context, attention_mask=mask)  
        encoder_out = output.last_hidden_state
        out, _ = self.lstm(encoder_out)
        out = self.dropout(out)
        out = self.fc_rnn(out[:, -1, :])  # 句子最后时刻的 hidden state
        return out